package com.pengpeng.zookeeper.lock;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @author 彭鹏
 * @date 2021/10/21.
 */
@Slf4j
public class DistributedLock {

    private String connectStr ="192.168.101.146:2181";
    private int timeOut = 5000;
    private ZooKeeper zkClient;
    private String rootNode = "locks";
    private String subNode = "seq-";
    private String waitPath;
    private CountDownLatch connectLatch = new CountDownLatch(1);
    private CountDownLatch waitLatch = new CountDownLatch(1);
    private String currentNode;

    public DistributedLock() throws KeeperException, InterruptedException, IOException {
        zkClient = new ZooKeeper(connectStr, timeOut, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    connectLatch.countDown();
                }
                if (event.getType() == Event.EventType.NodeDeleted && event.getPath().equals(waitPath)) {
                    waitLatch.countDown();
                }
            }
        });
        connectLatch.await();
        //获取根节点状态
        Stat exists = zkClient.exists("/" + rootNode, false);
        //如果根节点不存在，则创建永久根节点
        if (exists == null) {
            zkClient.create("/" + rootNode, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }

    public boolean zkLock() {
        try {
            currentNode = zkClient.create("/" + rootNode + "/" + subNode, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            List<String> childrenNodes = zkClient.getChildren("/" + rootNode, false);
            //列表只有一个子节点，说明client获得锁
            if (childrenNodes.size() == 1) {
                log.error("{}获取到锁",Thread.currentThread().getName());
                return true;
            } else {
                Collections.sort(childrenNodes);
                String thisNode = currentNode.substring(("/" + rootNode + "/").length());
                int index = childrenNodes.indexOf(thisNode);
                if (index == -1) {
                    //数据异常
                    log.error("数据异常");
                } else if (index == 0) {
                    //当前client获得锁
                    log.error("{}获取到锁",Thread.currentThread().getName());
                    return true;
                } else {
                    this.waitPath = "/" + rootNode + "/" + childrenNodes.get(index - 1);
                    zkClient.getData(waitPath, true, new Stat());
                    waitLatch.await();
                    log.error("{}获取到锁",Thread.currentThread().getName());
                    return true;
                }
            }
        } catch (Exception e) {
            log.error(Thread.currentThread().getName()+"获取锁异常："+e.getMessage());
        }
        return false;
    }

    public boolean zkUnlock() {
        try {
            zkClient.delete(this.currentNode, -1);
            log.error("{}释放锁",Thread.currentThread().getName());
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 核心线程池大小
     */
    public static final int CORE_POOL_SIZE = 5;
    /**
     * 最大线程池大小
     */
    public static final int MAX_POOL_SIZE = 10;
    /**
     * 阻塞任务队列大小
     */
    public static final int QUEUE_CAPACITY = 100;
    /**
     * 空闲线程存活时间
     */
    public static final Long KEEP_ALIVE_TIME = 1L;
    public static void main(String args[]) throws Exception {
        ThreadPoolExecutor threadPoolExecutor =
                new ThreadPoolExecutor(
                        CORE_POOL_SIZE,
                        MAX_POOL_SIZE,
                        KEEP_ALIVE_TIME,
                        TimeUnit.SECONDS,
                        new ArrayBlockingQueue<>(QUEUE_CAPACITY),
                        new ThreadPoolExecutor.CallerRunsPolicy());
        for (int i = 0;i<MAX_POOL_SIZE;i++){
            threadPoolExecutor.execute(()-> {
                DistributedLock lock = null;
                try {
                    lock = new DistributedLock();
                    lock.zkLock();
                    Thread.sleep(10);
                } catch (InterruptedException | KeeperException | IOException e) {
                    e.printStackTrace();
                } finally {
                    Objects.requireNonNull(lock).zkUnlock();
                }
            });
        }
        threadPoolExecutor.shutdown();
    }
}